package cn.ckx.coding.config;

import cn.ckx.coding.handler.AppAsyncUncaughtExceptionHandler;
import cn.ckx.coding.properties.AsyncConfigProperties;
import cn.hutool.core.thread.ThreadFactoryBuilder;
import org.springframework.aop.interceptor.AsyncUncaughtExceptionHandler;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableAsync;

import java.util.concurrent.*;

/**
 * AsyncConfig
 *
 * @author kunxin.cai
 */
@EnableAsync
@Configuration
public class AsyncConfig implements AsyncConfigurer {

    private static final String THREAD_PREFIX = "async-thread-";

    private final AsyncConfigProperties asyncConfigProperties;
    
    public AsyncConfig(AsyncConfigProperties asyncConfigProperties) {
        this.asyncConfigProperties = asyncConfigProperties;
    }
    
    @Override
    public Executor getAsyncExecutor() {
        return new IndexExecutor(asyncConfigProperties.getCorePoolSize(),
                asyncConfigProperties.getMaxPoolSize(),
                asyncConfigProperties.getKeepAliveTime(),
                TimeUnit.SECONDS,
                new ArrayBlockingQueue<>(asyncConfigProperties.getCapacity()),
                new ThreadPoolExecutor.AbortPolicy());
    }

    @Override
    public AsyncUncaughtExceptionHandler getAsyncUncaughtExceptionHandler() {
        return new AppAsyncUncaughtExceptionHandler();
    }

    private class IndexExecutor extends ThreadPoolExecutor{

        private int taskCount = 0;

        private int useSetCorePoolSize;

        private IndexExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, nameThreadFactory(), handler);
            this.useSetCorePoolSize=corePoolSize;
        }

        @Override
        public void execute(Runnable command) {
            synchronized (this){
                taskCount++;
                afterThreadIfNeed();
            }
            super.execute(command);
        }

        @Override
        public void setCorePoolSize(int corePoolSize) {
            if(corePoolSize>0) {
                //核心线程数大于活跃线程数
                if (corePoolSize > super.getActiveCount()) {
                    super.setCorePoolSize(corePoolSize);
                }
                //核心线程数大于最大线程数
                if (corePoolSize > super.getMaximumPoolSize()) {
                    super.setMaximumPoolSize(corePoolSize);
                }
                this.useSetCorePoolSize = corePoolSize;
            }
        }

        @Override
        protected void afterExecute(Runnable runnable, Throwable throwable) {
            try{
                super.afterExecute(runnable, throwable);
            }finally {
                synchronized (this){
                    taskCount--;
                    afterThreadIfNeed();
                }
            }
        }

        private void afterThreadIfNeed() {
            int needThreads=taskCount;
            //运行的线程数小于核心线程数
            if(needThreads<useSetCorePoolSize){
                needThreads=useSetCorePoolSize;
                //运行的线程数大于最大线程数
            }else if(needThreads>super.getMaximumPoolSize()){
                needThreads=super.getMaximumPoolSize();
            }
            super.setCorePoolSize(needThreads);
        }
    }

    private ThreadFactory nameThreadFactory() {
        return ThreadFactoryBuilder.create().setNamePrefix(THREAD_PREFIX).build();
    }
}
